#include <algorithm>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <map>
#include <random>

#include "api/graphwalker_basic_includes.hpp"

void read_edge_by_txt(std::string filename, Edge*& edge, eid_t& e_num)
{
    FILE * f = fopen(filename.c_str(), "r");
	assert(f != NULL);

	fseek(f, 0, SEEK_END);
    eid_t end = ftell(f);
	Edge temp;
	e_num = 0;
	fseek(f, 0, SEEK_SET);
	eid_t ft = ftell(f);
	while (ft < end)
	{
		if (read_edge_txt(f, &temp))
		{
			e_num++;
			// printf("\r%zu",e_num);
		}
		ft = ftell(f);
	}
	edge = new Edge[e_num];

	eid_t e_i = 0;
	fseek(f, 0, SEEK_SET);
	while ((eid_t)ftell(f) < end)
	{
		if (read_edge_txt(f, &temp))
		{
			edge[e_i] = temp;
			e_i++;
		}
	}
	fclose(f);

}

void convert_undirected(Edge*& read_edges, size_t& read_e_num)
{
    Edge* undirected_edges = new Edge[read_e_num * 2];
    for (eid_t e_i = 0; e_i < read_e_num; e_i++)
    {
        undirected_edges[e_i * 2] = read_edges[e_i];
        std::swap(read_edges[e_i].src, read_edges[e_i].dst);
        undirected_edges[e_i * 2 + 1] = read_edges[e_i];
    }
    delete[] read_edges;
    read_edges = undirected_edges;
    read_e_num *= 2;
}

void get_degree( Edge* edges, eid_t e_num, vid_t v_num, std::vector<vid_t>& degree)
{
    degree.resize(v_num);

    for(eid_t e_i = 0; e_i < e_num; e_i++) {
        degree[edges[e_i].src]++;
    }
}

void get_degree_directed( Edge* edges, eid_t e_num, vid_t v_num, std::vector<vid_t>& degree)
{
    degree.resize(v_num);

    for(eid_t e_i = 0; e_i < e_num; e_i++) {
        degree[edges[e_i].src]++;
    }  
}

void build_edge_container(Edge* edges, eid_t e_num, vid_t v_num, EdgeContainer* ec, const std::vector<vid_t>& vertex_out_degree)
{
    ec->adj_lists = new AdjList[v_num];
    ec->adj_units = new AdjUnit[e_num];
    eid_t chunk_edge_idx = 0;
    for(vid_t v_i = 0; v_i < v_num; v_i++)
    {
        ec->adj_lists[v_i].begin = ec->adj_units + chunk_edge_idx;
        ec->adj_lists[v_i].end = ec->adj_lists[v_i].begin;
        chunk_edge_idx += vertex_out_degree[v_i];
    }
    for(eid_t e_i = 0; e_i < e_num; e_i++)
    {
        auto e = edges[e_i];
        auto ep = ec->adj_lists[e.src].end++;
        ep->neighbor = e.dst;
    }
}

void sort_csr(EdgeContainer* csr, vid_t v_num)
{
    for (vid_t v_i = 0; v_i < v_num; v_i++) {
        std::sort(csr->adj_lists[v_i].begin, csr->adj_lists[v_i].end,
            [](const AdjUnit a, const AdjUnit b) {return a.neighbor < b.neighbor; });
    }
}

void DFS_degree(EdgeContainer* graph_csr, std::queue<vid_t>& stream, vid_t v_num, std::vector<vid_t>& graph_degree)
{
    std::vector<bool> is_visited(v_num, false);
    std::stack<vid_t> sta;
    vid_t max_degree = 0;
    vid_t max_d_i;

    for(vid_t v_i = 0; v_i < v_num; v_i++)
    {
        if(graph_degree[v_i] > max_degree && is_visited[v_i] == false){
            max_degree = graph_degree[v_i];
            max_d_i = v_i;
        }
    }
    sta.push(max_d_i);
    is_visited[max_d_i] = true;

    while(!sta.empty())
    {
        vid_t sta_top = sta.top();
        stream.push(sta_top);
        sta.pop();
        std::vector<vid_t> t;
        for(auto a = graph_csr->adj_lists[sta_top].begin; a < graph_csr->adj_lists[sta_top].end; a++) {
            if(is_visited[a->neighbor] == false){
                t.push_back(a->neighbor);
                is_visited[a->neighbor] = true;
            }
        }
        sort(t.begin(),t.end(),[&](vid_t a, vid_t b){
            return graph_degree[a] < graph_degree[b];
        });
        for(vid_t i =0; i< t.size();i++)
        {
            sta.push(t[i]);
        }
    }
    printf("stream size: %zu    %.2f%%\n",stream.size(),stream.size()/(float)v_num * 100);
    for(vid_t v = 0; v < v_num; v++){
        if(is_visited[v] == false){
            stream.push(v);
            is_visited[v] = true;
        }
    }
    assert(v_num == stream.size());
}

void sequence(std::queue<vid_t>& stream, vid_t v_num)
{
    for (vid_t v = 0; v < v_num; v++)
        stream.push(v);
}

void random_stream(std::queue<vid_t>& stream, vid_t v_num)
{
    // 没有与参数列表匹配的shuffle实例
    std::vector<vid_t> arr(v_num);
    for (vid_t v = 0; v < v_num; v++)
        arr[v] = v;
    
    // 使用random_device生成随机数种子
    std::random_device rd;
    // 使用默认随机数引擎和种子
    std::default_random_engine rng(rd());
    
    // 使用shuffle函数将数组打乱
    std::shuffle(arr.begin(), arr.end(), rng); 
    for (vid_t v = 0; v < v_num; v++)
    {
        stream.push(arr[v]);
        // std::cout << arr[v] << " ";
    }     
    //std::cout << std::endl;
}

void vertex_to_partition(std::string filename, std::string partition_mode, std::vector<std::vector<vid_t>> partition_vertex_array, std::vector<vid_t>& ver2part)
{
    mkdir((filename+"_ver2part/").c_str(), 0777);
    for(size_t p_i = 0; p_i < partition_vertex_array.size(); p_i++)
    {
        for(vid_t v_i = 0; v_i < partition_vertex_array[p_i].size(); v_i++)
        {
            ver2part[partition_vertex_array[p_i][v_i]] = p_i;
        }
    }
    std::ofstream ver2part_file(filename + "_" + partition_mode + "_" + std::to_string(partition_vertex_array.size()) + ".ver2part");
    for(std::vector<uint32_t>::iterator it = ver2part.begin(); it != ver2part.end(); ++it) {
        ver2part_file << *it << std::endl;
    }
    ver2part_file.close();
}

void outPLEL(EdgeContainer* graph_csr, vid_t v_num, std::vector<vid_t>& degree, std::vector<vid_t>& outPL, std::vector<eid_t>& outEL)
{
    vid_t outid = 0;
    eid_t out_curnum = 0;

    for(vid_t i = 0; i < v_num; i++)
    {
        outid += degree[i];
        outPL[i + 1] = outid;
        AdjList* v1 = &graph_csr->adj_lists[i];
        for (vid_t j = 0; j < degree[i]; j++)
        {
            outEL[out_curnum] = (v1->begin + j)->neighbor; // outAdjList
            out_curnum += 1;
        }
    }
}

void partition_relabel(std::string filename, vid_t v_num, int partition_num, std::string partition_mode, std::string stream_mode)
{
    logstream(LOG_DEBUG) << "Start to partitoning." << std::endl;

    Edge* graph_edges;
    eid_t e_num;

    read_edge_by_txt(filename, graph_edges, e_num);
    logstream(LOG_INFO) << "Read Graph." << std::endl;

    // convert_undirected(graph_edges, e_num);
    // logstream(LOG_INFO) << "Convert_undirecteds finished." << std::endl;
    logstream(LOG_INFO) << "enum: " << e_num << std::endl;

    std::vector<vid_t> graph_degree;
    
    get_degree(graph_edges, e_num, v_num, graph_degree);
    logstream(LOG_INFO) << "Get_degree finished." << std::endl;

    EdgeContainer* graph_csr = new EdgeContainer();
    build_edge_container(graph_edges, e_num, v_num, graph_csr, graph_degree);
    sort_csr(graph_csr, v_num);
    logstream(LOG_INFO) << "Build and sort CSR finished." << std::endl;

    std::vector<vid_t> outPL(v_num + 1);
    std::vector<eid_t> outEL(e_num);
    outPLEL(graph_csr, v_num, graph_degree, outPL, outEL);
    
    float beta = 1.05;

    std::vector<std::vector<vid_t>>partition_vertex_array(partition_num);
    if(partition_mode == "naive") // 按顶点id升序进行图划分
    {
        Naive partitioner;
        partitioner.naive(v_num, e_num, partition_vertex_array, beta, outPL);
    }
    else if(partition_mode == "MPGP")
    {
        MPGP partitioner;
        partitioner.mpgp(graph_csr, v_num, partition_vertex_array, beta, graph_degree, stream_mode);
    }     
    else if(partition_mode == "AGP")
    {
        AGPFlag *partitioner = new AGPFlag(v_num, partition_num);
        // AGPVec *partitioner = new AGPVec(v_num, partition_num);
        // AGPMap *partitioner = new AGPMap(v_num, partition_num);
        // AGPSet *partitioner = new AGPSet(v_num, partition_num);
        // partitioner->testVec();
        // exit(-1);
        partitioner->agp(graph_csr, v_num, partition_vertex_array, beta, graph_degree, outPL, outEL, stream_mode);
    }          
    std::vector<vid_t> ver2part(v_num);
    vertex_to_partition(filename, partition_mode, partition_vertex_array, ver2part);
}


int main(int argc, const char ** argv) {
    set_argc(argc, argv);
    metrics m("mpgp");

    std::string filename = get_option_string("file", "./dataset/karate.edgelist");
    vid_t N = get_option_int("N", 34); // Number of vertices
    int P = get_option_int("P", 3); // Number of partitions
    std::string partition_mode = get_option_string("partition_mode", "MPGP");
    std::string stream_mode = get_option_string("stream_mode", "DFS_degree");
    

    assert(N > 0 && P > 0);
    logstream(LOG_DEBUG) << "N, P, steam_mode: " << N << " " << P << " " << stream_mode << std::endl;

    partition_relabel(filename, N, P, partition_mode, stream_mode);
}